from __future__ import annotations

import json
import sys
from typing import List

from adapter.diting_optimization import (
    DatasetItem,
    StageConfig,
    build_callback_config,
    build_optimization_payload,
)
from rag_service import RAGService


def build_sample_dataset() -> List[DatasetItem]:
    """Create placeholder evaluation samples to demonstrate payload structure."""
    return [
        DatasetItem(
            id="sample-001",
            user_input="diting RAG 系统的主要功能是什么？",
            expected_output="diting RAG 系统用于构建检索增强生成（RAG）工作流。",
        )
        # 减少测试数据集大小以减少优化时间
    ]


def main() -> int:
    """Invoke diting optimization endpoint and print best parameters."""
    service = RAGService()
    if not service.settings.optimization_base_url:
        print("RAG_OPTIMIZATION_BASE_URL 未配置，无法发起优化请求。", file=sys.stderr)
        return 1

    stage = StageConfig(
        name="retrieval.semantic_search",
        parameters={"search_k": service.settings.search_k},
        similarity_threshold_range=(0.7, 0.9),  # 缩小范围以减少优化时间
        similarity_threshold_step=0.1,          # 保持步长
        context_recall_tokens_range=(512, 1024), # 缩小范围以减少优化时间
        context_recall_tokens_step=256,
    )

    callback_cfg = build_callback_config()  # 使用默认不回调；可按需传入 URL。

    # 构建 global_context，包含 remote_base_url 用于检索阶段
    # 这里应该设置为 RAG 服务的 URL，而不是 LLM API 的 URL
    global_ctx = {
        "remote_base_url": "http://localhost:8000",  # 假设 RAG 服务运行在 8000 端口
    }

    payload = build_optimization_payload(
        schema_id="rag_default_schema",
        run_id="rag-opt-demo",
        stage_sequence=["retrieval.semantic_search"],
        stages=[stage],
        dataset_items=build_sample_dataset(),
        metric={"type": "average_similarity"},
        global_context=global_ctx,
        callback=callback_cfg,
    )

    try:
        best_params, response = service.optimize_pipeline(payload)
    except RuntimeError as exc:
        print(f"优化请求失败: {exc}", file=sys.stderr)
        return 1

    print("建议的最优参数：")
    print(json.dumps(best_params, indent=2, ensure_ascii=False))
    print("\n完整响应数据：")
    print(json.dumps(response, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    exit_code = main()
    sys.exit(exit_code)
